{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## dbSNP\n",
    "\n",
    "Use to create Hail Tables for dbSNP, after downloading raw data from https://ftp.ncbi.nih.gov/snp/. \n",
    "\n",
    "Raw data downloaded with Hail Batch, see `hail/datasets/extract/extract_dbSNP.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import hail as hl\n",
    "hl.init()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Hail Tables from GRCh37 and GRCh38 assembly reports\n",
    "\n",
    "The contigs in the VCFs are [RefSeq](https://www.ncbi.nlm.nih.gov/refseq/) accession numbers, and need to be mapped back to the appropriate chromosome for each reference genome.\n",
    "\n",
    "The GRCh37 assembly can be found [here](https://www.ncbi.nlm.nih.gov/assembly/GCF_000001405.25), and the assembly report mapping chromosomes to RefSeq sequences can be found [here](https://ftp.ncbi.nlm.nih.gov/genomes/all/GCF/000/001/405/GCF_000001405.25_GRCh37.p13/GCF_000001405.25_GRCh37.p13_assembly_report.txt).\n",
    "\n",
    "The GRCh38 assembly can be found [here](https://www.ncbi.nlm.nih.gov/assembly/GCF_000001405.39), and the assembly report mapping chromosomes to RefSeq sequences can be found [here](https://ftp.ncbi.nlm.nih.gov/genomes/all/GCF/000/001/405/GCF_000001405.39_GRCh38.p13/GCF_000001405.39_GRCh38.p13_assembly_report.txt)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### GRCh37"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "ht = hl.import_table(\"gs://hail-datasets-tmp/dbSNP/GCF_000001405.25_GRCh37.p13_assembly_report.txt\", \n",
    "                     no_header=True, \n",
    "                     comment=\"#\",\n",
    "                     delimiter=\"\\t\", \n",
    "                     missing=\"na\")\n",
    "\n",
    "field_names = ['sequence_name','sequence_role','assigned_molecule',\n",
    "               'assigned_molecule_location/type', 'genbank_accn', 'relationship', \n",
    "               'refseq_accn', 'assembly_unit', 'sequence_length', 'ucsc_style_name']\n",
    "\n",
    "name = \"dbSNP\"\n",
    "version = \"154\"\n",
    "build = \"GRCh37\"\n",
    "n_rows = ht.count()\n",
    "n_partitions = ht.n_partitions()\n",
    "\n",
    "ht = ht.annotate_globals(\n",
    "    metadata=hl.struct(\n",
    "        name=name,\n",
    "        version=version,\n",
    "        reference_genome=build,\n",
    "        n_rows=n_rows,\n",
    "        n_partitions=n_partitions\n",
    "    )\n",
    ")\n",
    "ht = ht.rename(dict(zip([f\"f{i}\" for i in range(10)], field_names)))\n",
    "ht = ht.drop(\"relationship\").key_by(\"refseq_accn\")\n",
    "\n",
    "ht.write(\"gs://hail-datasets-us/NCBI_assembly_report_p13_GRCh37.ht\")\n",
    "ht = hl.read_table(\"gs://hail-datasets-us/NCBI_assembly_report_p13_GRCh37.ht\")\n",
    "ht.describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### GRCh38"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "ht = hl.import_table(\"gs://hail-datasets-tmp/dbSNP/GCF_000001405.39_GRCh38.p13_assembly_report.txt\", \n",
    "                     no_header=True, \n",
    "                     comment=\"#\",\n",
    "                     delimiter=\"\\t\", \n",
    "                     missing=\"na\")\n",
    "\n",
    "field_names = ['sequence_name','sequence_role','assigned_molecule',\n",
    "               'assigned_molecule_location/type', 'genbank_accn', 'relationship', \n",
    "               'refseq_accn', 'assembly_unit', 'sequence_length', 'ucsc_style_name']\n",
    "\n",
    "name = \"dbSNP\"\n",
    "version = \"154\"\n",
    "build = \"GRCh38\"\n",
    "n_rows = ht.count()\n",
    "n_partitions = ht.n_partitions()\n",
    "\n",
    "ht = ht.annotate_globals(\n",
    "    metadata=hl.struct(\n",
    "        name=name,\n",
    "        version=version,\n",
    "        reference_genome=build,\n",
    "        n_rows=n_rows,\n",
    "        n_partitions=n_partitions\n",
    "    )\n",
    ")\n",
    "ht = ht.rename(dict(zip([f\"f{i}\" for i in range(10)], field_names)))\n",
    "ht = ht.drop(\"relationship\").key_by(\"refseq_accn\")\n",
    "\n",
    "ht.write(\"gs://hail-datasets-us/NCBI_assembly_report_p13_GRCh38.ht\")\n",
    "ht = hl.read_table(\"gs://hail-datasets-us/NCBI_assembly_report_p13_GRCh38.ht\")\n",
    "ht.describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Hail Tables for dbSNP\n",
    "\n",
    "Now we can use the assembly report for each reference genome build to map from RefSeq accession numbers to chromosomes, and create Hail Tables. There are no samples or entries in the dbSNP VCFs. Some helpful information about the dbSNP VCFs is available [here](https://www.ncbi.nlm.nih.gov/snp/docs/products/vcf/redesign/).\n",
    "\n",
    "We will create two Hail Tables for each reference genome build, both keyed by `[\"locus\", \"alleles\"]`:\n",
    "\n",
    "  - Table with all fields from the imported VCF (e.g. `gs://hail-datasets-us/dbSNP_154_GRCh37.ht`)\n",
    "  - Table with only the rsID field (e.g. `gs://hail-datasets-us/dbSNP_rsid_154_GRCh37.ht`)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First load VCFs to get all the contigs present in each dataset so we can create a mapping to used to recode contigs from RefSeq accession numbers to GRCh37/38 builds. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mt37 = hl.import_vcf(f\"gs://hail-datasets-tmp/dbSNP/dbSNP_154_GRCh37.vcf.bgz\", \n",
    "                     header_file=f\"gs://hail-datasets-tmp/dbSNP/dbSNP_154_GRCh37_header_only.vcf.txt\", \n",
    "                     reference_genome=None, \n",
    "                     min_partitions=512)\n",
    "\n",
    "mt38 = hl.import_vcf(f\"gs://hail-datasets-tmp/dbSNP/dbSNP_154_GRCh38.vcf.bgz\", \n",
    "                     header_file=f\"gs://hail-datasets-tmp/dbSNP/dbSNP_154_GRCh38_header_only.vcf.txt\", \n",
    "                     reference_genome=None, \n",
    "                     min_partitions=512)\n",
    "\n",
    "mt37 = mt37.checkpoint(f\"gs://hail-datasets-tmp/checkpoints/dbSNP_154_GRCh37_no_coding.mt\", \n",
    "                       _read_if_exists=True, \n",
    "                       overwrite=False)\n",
    "\n",
    "mt38 = mt38.checkpoint(f\"gs://hail-datasets-tmp/checkpoints/dbSNP_154_GRCh38_no_coding.mt\", \n",
    "                       _read_if_exists=True, \n",
    "                       overwrite=False)\n",
    "\n",
    "# To get all contigs present for recoding to correct reference genome mapping\n",
    "contigs_present37 = mt37.aggregate_rows(hl.agg.collect_as_set(mt37.locus.contig))\n",
    "contigs_present38 = mt38.aggregate_rows(hl.agg.collect_as_set(mt38.locus.contig))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Load NCBI assembly reports with RefSeq mappings\n",
    "assembly37_ht = hl.read_table(\"gs://hail-datasets-us/NCBI_assembly_report_p13_GRCh37.ht\")\n",
    "assembly37_ht = assembly37_ht.annotate(\n",
    "    contig = hl.if_else(assembly37_ht.sequence_role == \"unlocalized-scaffold\", \n",
    "                        assembly37_ht.genbank_accn, \n",
    "                        assembly37_ht.assigned_molecule)\n",
    ")\n",
    "assembly38_ht = hl.read_table(\"gs://hail-datasets-us/NCBI_assembly_report_p13_GRCh38.ht\")\n",
    "\n",
    "# Map RefSeq identifiers to chromosomes for GRCh37 using the \"contig\" field we created in assembly report\n",
    "rg37 = hl.get_reference(\"GRCh37\")\n",
    "refseq_to_chr37 = dict(zip(assembly37_ht.refseq_accn.collect(), assembly37_ht.contig.collect()))\n",
    "refseq_to_chr37 = {k: v for k, v in refseq_to_chr37.items() if k in contigs_present37 and v in rg37.contigs}\n",
    "\n",
    "# Map RefSeq identifiers to chromosomes for GRCh38 using the \"ucsc_style_name\" field in assembly report\n",
    "rg38 = hl.get_reference(\"GRCh38\")\n",
    "refseq_to_chr38 = dict(zip(assembly38_ht.refseq_accn.collect(), assembly38_ht.ucsc_style_name.collect()))\n",
    "refseq_to_chr38 = {k: v for k, v in refseq_to_chr38.items() if k in contigs_present38 and v in rg38.contigs}\n",
    "\n",
    "recodings = {\n",
    "    \"GRCh37\": refseq_to_chr37, \n",
    "    \"GRCh38\": refseq_to_chr38\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the function and known keys below to convert allele frequency arrays to structs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert array of strings like hl.array([\"GnomAD:.,1,3.187e-05\",\"TOPMED:.,1,2.389e-05\"]) to a struct\n",
    "def arr_str_to_struct(hl_array, known_keys):\n",
    "    _dict = hl.dict(\n",
    "        hl_array.map(\n",
    "            lambda x: (\"_\" + x.split(\":\")[0], \n",
    "                       x.split(\":\")[1].split(\",\").map(lambda x: hl.if_else(x == \".\", hl.missing(hl.tfloat), hl.float(x))))\n",
    "        )\n",
    "    )\n",
    "    _struct =  hl.rbind(_dict, lambda d: hl.struct(**{k: _dict.get(k) for k in known_keys}))\n",
    "    return _struct\n",
    "\n",
    "# To get all possible keys for allele frequency arrays after loading VCF as MatrixTable\n",
    "#     known_keys_FREQ = mt.aggregate_rows(\n",
    "#         hl.agg.explode(\n",
    "#             lambda x: hl.agg.collect_as_set(x), mt.info.FREQ.split(\"\\\\|\").map(lambda x: x.split(\":\")[0])\n",
    "#         )\n",
    "#     )\n",
    "\n",
    "known_keys = ['GENOME_DK','TWINSUK','dbGaP_PopFreq','Siberian','Chileans',\n",
    "              'FINRISK','HapMap','Estonian','ALSPAC','GoESP',\n",
    "              'TOPMED','PAGE_STUDY','1000Genomes','Korea1K','ChromosomeY',\n",
    "              'ExAC','Qatari','GoNL','MGP','GnomAD',\n",
    "              'Vietnamese','GnomAD_exomes','PharmGKB','KOREAN','Daghestan',\n",
    "              'HGDP_Stanford','NorthernSweden','SGDP_PRJ']\n",
    "known_keys_FREQ = list(map(lambda x: f\"_{x}\", known_keys))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now can read in VCF files again as MatrixTables with the correct contig recodings, and reformat the allele frequency information in `info.FREQ` and the clinical attributes in `info`.\n",
    "\n",
    "Note that we are specifying a separate header file in the `hl.import_vcf` calls in the cell below. \n",
    "\n",
    "To make parsing strings easier, the following INFO fields in the VCF headers were changed from `Number=.` to `Number=1`: FREQ, CLNHGVS, CLNVI, CLNORIGIN, CLNSIG, CLNDISB, CLNDN, CLNREVSTAT, CLNACC. \n",
    "\n",
    "The modified VCF headers used are available in `gs://hail-datasets-tmp/dbSNP`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "name = \"dbSNP\"\n",
    "version = \"154\"\n",
    "builds = [\"GRCh37\", \"GRCh38\"]\n",
    "\n",
    "for build in builds:\n",
    "    mt = hl.import_vcf(f\"gs://hail-datasets-tmp/{name}/{name}_{version}_{build}.vcf.bgz\", \n",
    "                       header_file=f\"gs://hail-datasets-tmp/{name}/{name}_{version}_{build}_header_only.vcf.txt\", \n",
    "                       contig_recoding=recodings[build], \n",
    "                       skip_invalid_loci=True, \n",
    "                       reference_genome=build, \n",
    "                       min_partitions=512)\n",
    "\n",
    "    # First annotation, go from str to array<str> for FREQ\n",
    "    mt = mt.annotate_rows(\n",
    "        info = mt.info.annotate(\n",
    "            FREQ = hl.or_missing(hl.is_defined(mt.info.FREQ), mt.info.FREQ.split(\"\\\\|\"))\n",
    "        )\n",
    "    )\n",
    "    # Second annotation, turn array<str> into a struct for FREQ\n",
    "    mt = mt.annotate_rows(\n",
    "        info = mt.info.annotate(\n",
    "            FREQ = hl.or_missing(hl.is_defined(mt.info.FREQ), \n",
    "                                 arr_str_to_struct(mt.info.FREQ, known_keys_FREQ))\n",
    "        )\n",
    "    )\n",
    "    # Reformat clinical attributes from str to array\n",
    "    mt = mt.annotate_rows(\n",
    "        info = mt.info.annotate(\n",
    "            CLNHGVS = hl.or_missing(\n",
    "                hl.is_defined(mt.info.CLNHGVS), \n",
    "                mt.info.CLNHGVS.split(\"(?:(\\|)|(\\,))\")).map(lambda x: hl.if_else((x == \".\"), hl.missing(hl.tstr), x)),\n",
    "            CLNVI = hl.or_missing(\n",
    "                hl.is_defined(mt.info.CLNVI), \n",
    "                mt.info.CLNVI.split(\"(?:(\\|)|(\\,))\")).filter(lambda x: x != \".\"),\n",
    "            CLNORIGIN = hl.or_missing(\n",
    "                hl.is_defined(mt.info.CLNORIGIN), \n",
    "                mt.info.CLNORIGIN.split(\"(?:(\\|)|(\\,))\")).filter(lambda x: x != \".\"),\n",
    "            CLNSIG = hl.or_missing(\n",
    "                hl.is_defined(mt.info.CLNSIG), \n",
    "                mt.info.CLNSIG.split(\"(?:(\\|)|(\\,))\")).filter(lambda x: x != \".\"),\n",
    "            CLNDISDB = hl.or_missing(\n",
    "                hl.is_defined(mt.info.CLNDISDB), \n",
    "                mt.info.CLNDISDB.split(\"(?:(\\|)|(\\,))\")).filter(lambda x: x != \".\"),\n",
    "            CLNDN = hl.or_missing(\n",
    "                hl.is_defined(mt.info.CLNDN), \n",
    "                mt.info.CLNDN.split(\"(?:(\\|)|(\\,))\")).filter(lambda x: x != \".\"),\n",
    "            CLNREVSTAT = hl.or_missing(\n",
    "                hl.is_defined(mt.info.CLNREVSTAT), \n",
    "                mt.info.CLNREVSTAT.split(\"(?:(\\|)|(\\,))\")).filter(lambda x: x != \".\"),\n",
    "            CLNACC = hl.or_missing(\n",
    "                hl.is_defined(mt.info.CLNACC), \n",
    "                mt.info.CLNACC.split(\"(?:(\\|)|(\\,))\")).filter(lambda x: x != \".\")\n",
    "        )\n",
    "    )\n",
    "    \n",
    "    mt = mt.checkpoint(f\"gs://hail-datasets-tmp/checkpoints/{name}_{version}_{build}.mt\", \n",
    "                       _read_if_exists=True, \n",
    "                       overwrite=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we can just grab the `rows` table since we have no sample or entry information in the MatrixTable. \n",
    "\n",
    "From there, we need to filter the biallelic and multiallelic variants into separate tables, split the multiallelic variants, and then union the split multiallelic table rows back with the biallelic table rows.\n",
    "\n",
    "The allele frequency arrays start with the reference allele which is then followed by alternate alleles as ordered in the ALT column (from the VCF). So after splitting we can index the array with `a_index` to pull out the relevant allele frequency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "name = \"dbSNP\"\n",
    "version = \"154\"\n",
    "builds = [\"GRCh37\", \"GRCh38\"]\n",
    "\n",
    "for build in builds:\n",
    "    # No samples or entries in MT, just grab table with the rows\n",
    "    mt = hl.read_matrix_table(f\"gs://hail-datasets-tmp/checkpoints/{name}_{version}_{build}.mt\")\n",
    "    ht = mt.rows()\n",
    "       \n",
    "    ht_ba = ht.filter(hl.len(ht.alleles) <= 2)\n",
    "    ht_ba = ht_ba.checkpoint(f\"gs://hail-datasets-tmp/checkpoints/{name}_{version}_{build}_biallelic.ht\", \n",
    "                             _read_if_exists=True, \n",
    "                             overwrite=False)\n",
    "\n",
    "    ht_ma = ht.filter(hl.len(ht.alleles) > 2)\n",
    "    ht_ma = ht_ma.checkpoint(f\"gs://hail-datasets-tmp/checkpoints/{name}_{version}_{build}_multiallelic.ht\", \n",
    "                             _read_if_exists=True, \n",
    "                             overwrite=False)\n",
    "\n",
    "    ht_split = hl.split_multi(ht_ma, keep_star=True, permit_shuffle=True)\n",
    "    ht_split = ht_split.repartition(64, shuffle=False)\n",
    "    ht_split = ht_split.checkpoint(f\"gs://hail-datasets-tmp/checkpoints/{name}_{version}_{build}_split_multiallelic.ht\", \n",
    "                                   _read_if_exists=True, \n",
    "                                   overwrite=False)\n",
    "    \n",
    "    # Next, have to fix indices and union ht_split with ht_ba\n",
    "    ht_union = ht_ba.union(ht_split, unify=True)\n",
    "    ht_union = ht_union.annotate(\n",
    "        a_index = hl.if_else(hl.is_missing(ht_union.a_index), 1, ht_union.a_index),\n",
    "        was_split = hl.if_else(hl.is_missing(ht_union.was_split), False, ht_union.was_split),\n",
    "        old_locus = hl.if_else(hl.is_missing(ht_union.old_locus), ht_union.locus, ht_union.old_locus),\n",
    "        old_alleles = hl.if_else(hl.is_missing(ht_union.old_alleles), ht_union.alleles, ht_union.old_alleles)\n",
    "    )\n",
    "    ht_union = ht_union.checkpoint(f\"gs://hail-datasets-tmp/checkpoints/{name}_{version}_{build}_unioned.ht\", \n",
    "                                   _read_if_exists=True, \n",
    "                                   overwrite=False)\n",
    "    \n",
    "    # Arrays for AFs start w/ ref allele in index 0, so just use a_index to get alternate AFs\n",
    "    ht = ht_union.annotate(\n",
    "        info = ht_union.info.annotate(\n",
    "            FREQ = ht_union.info.FREQ.annotate(\n",
    "                **{k: hl.or_missing(hl.is_defined(ht_union.info.FREQ[k]), \n",
    "                                    ht_union.info.FREQ[k][ht_union.a_index]) \n",
    "                   for k in known_keys_FREQ}\n",
    "            )\n",
    "        )\n",
    "    )\n",
    "    ht = ht.repartition(512, shuffle=True)\n",
    "    ht = ht.checkpoint(f\"gs://hail-datasets-tmp/checkpoints/{name}_{version}_{build}.ht\", \n",
    "                       _read_if_exists=True, \n",
    "                       overwrite=False)\n",
    "\n",
    "    n_rows = ht.count()\n",
    "    n_partitions = ht.n_partitions()\n",
    "\n",
    "    ht = ht.annotate_globals(\n",
    "        metadata=hl.struct(\n",
    "            name=name,\n",
    "            version=version,\n",
    "            reference_genome=build,\n",
    "            n_rows=n_rows,\n",
    "            n_partitions=n_partitions\n",
    "        )\n",
    "    )\n",
    "    ht.write(f\"gs://hail-datasets-us/{name}_{version}_{build}.ht\")\n",
    "    ht = hl.read_table(f\"gs://hail-datasets-us/{name}_{version}_{build}.ht\")\n",
    "    ht.describe()\n",
    "    print(str(hl.eval(ht.metadata)) + \"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Also write tables with only the rsID field, for smaller tables that just map `[locus, alleles]` to `rsID`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "name = \"dbSNP\"\n",
    "version = \"154\"\n",
    "builds = [\"GRCh37\", \"GRCh38\"]\n",
    "\n",
    "for build in builds:\n",
    "    # Write table with only rsid's\n",
    "    ht_rsid = hl.read_table(f\"gs://hail-datasets-us/{name}_{version}_{build}.ht\")\n",
    "    ht_rsid = ht_rsid.select(\"rsid\")\n",
    "\n",
    "    n_rows = ht_rsid.count()\n",
    "    n_partitions = ht_rsid.n_partitions()\n",
    "\n",
    "    ht_rsid = ht_rsid.annotate_globals(\n",
    "        metadata=hl.struct(\n",
    "            name=f\"{name}_rsid\",\n",
    "            version=version,\n",
    "            reference_genome=build,\n",
    "            n_rows=n_rows,\n",
    "            n_partitions=n_partitions\n",
    "        )\n",
    "    )\n",
    "    ht_rsid.write(f\"gs://hail-datasets-us/{name}_rsid_{version}_{build}.ht\")\n",
    "    ht_rsid = hl.read_table(f\"gs://hail-datasets-us/{name}_rsid_{version}_{build}.ht\")\n",
    "    ht_rsid.describe()\n",
    "    print(str(hl.eval(ht_rsid.metadata)) + \"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To check uniqueness of keys\n",
    "tables = [\"gs://hail-datasets-us/dbSNP_rsid_154_GRCh37.ht\", \"gs://hail-datasets-us/dbSNP_rsid_154_GRCh38.ht\"]\n",
    "for table in tables:\n",
    "    ht = hl.read_table(table)\n",
    "    \n",
    "    ht_count = ht.count()\n",
    "    print(f\"n = {ht_count}\")\n",
    "    ht_distinct_count = ht.distinct().count()\n",
    "    print(f\"n_distinct = {ht_distinct_count}\")\n",
    "    \n",
    "    if ht_count == ht_distinct_count:\n",
    "        print(f\"{table} rows unique\\n\")\n",
    "    else:\n",
    "        print(f\"{table} rows NOT unique\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add dbSNP to datasets API and annotation DB\n",
    "\n",
    "Now we can add the tables we created to `hail/python/hail/experimental/datasets.json`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "datasets_path = os.path.abspath(\"../../hail/python/hail/experimental/datasets.json\")\n",
    "with open(datasets_path, \"r\") as f:\n",
    "    datasets = json.load(f)\n",
    "\n",
    "names = [\"dbSNP\", \"dbSNP_rsid\"]\n",
    "version = \"154\"\n",
    "builds = [\"GRCh37\", \"GRCh38\"]\n",
    "\n",
    "gcs_us_url_root = \"gs://hail-datasets-us\"\n",
    "gcs_eu_url_root = \"gs://hail-datasets-eu\"\n",
    "aws_us_url_root = \"s3://hail-datasets-us-east-1\"\n",
    "\n",
    "for name in names:\n",
    "    json_entry = {\n",
    "        \"annotation_db\": {\n",
    "            \"key_properties\": []\n",
    "        },\n",
    "        \"description\": \"dbSNP: Reference SNP (rs or RefSNP) Hail Table. The database includes both common and rare single-base nucleotide variation (SNV), short (=< 50bp) deletion/insertion polymorphisms, and other classes of small genetic variations.\",\n",
    "        \"url\": \"https://www.ncbi.nlm.nih.gov/snp/docs/RefSNP_about/\",\n",
    "        \"versions\": [\n",
    "            {\n",
    "                \"reference_genome\": builds[0],\n",
    "                \"url\": {\n",
    "                    \"aws\": {\n",
    "                        \"us\": f\"{aws_us_url_root}/{name}_{version}_{builds[0]}.ht\"\n",
    "                    },\n",
    "                    \"gcp\": {\n",
    "                        \"eu\": f\"{gcs_eu_url_root}/{name}_{version}_{builds[0]}.ht\",\n",
    "                        \"us\": f\"{gcs_us_url_root}/{name}_{version}_{builds[0]}.ht\"\n",
    "                    }\n",
    "                },\n",
    "                \"version\": version\n",
    "            },\n",
    "            {\n",
    "                \"reference_genome\": builds[1],\n",
    "                \"url\": {\n",
    "                    \"aws\": {\n",
    "                        \"us\": f\"{aws_us_url_root}/{name}_{version}_{builds[1]}.ht\"\n",
    "                    },\n",
    "                    \"gcp\": {\n",
    "                        \"eu\": f\"{gcs_eu_url_root}/{name}_{version}_{builds[1]}.ht\",\n",
    "                        \"us\": f\"{gcs_us_url_root}/{name}_{version}_{builds[1]}.ht\"\n",
    "                    }\n",
    "                },\n",
    "                \"version\": version\n",
    "            }            \n",
    "        ]\n",
    "    }\n",
    "    \n",
    "    if name == \"dbSNP_rsid\":\n",
    "        json_entry[\"description\"] = \"dbSNP: This Hail Table contains a mapping from locus/allele pairs to Reference SNP IDs (rsID). For the full dataset, see dbSNP.\"\n",
    "    \n",
    "    datasets[name] = json_entry\n",
    "\n",
    "# Write new entries back to datasets.json config:\n",
    "with open(datasets_path, \"w\") as f:\n",
    "    json.dump(datasets, f, sort_keys=True, ensure_ascii=False, indent=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After adding tables to `datasets.json`, create .rst schema files for docs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Create/update schema .rst file\n",
    "import textwrap\n",
    "\n",
    "output_dir = os.path.abspath(\"../../hail/python/hail/docs/datasets/schemas\")\n",
    "datasets_path = os.path.abspath(\"../../hail/python/hail/experimental/datasets.json\")\n",
    "with open(datasets_path, \"r\") as f:\n",
    "    datasets = json.load(f)\n",
    "\n",
    "names = [\"dbSNP\", \"dbSNP_rsid\"]\n",
    "for name in names:\n",
    "    versions = sorted(set(dataset[\"version\"] for dataset in datasets[name][\"versions\"]))\n",
    "    if not versions:\n",
    "        versions = [None]\n",
    "    reference_genomes = sorted(set(dataset[\"reference_genome\"] for dataset in datasets[name][\"versions\"]))\n",
    "    if not reference_genomes:\n",
    "        reference_genomes = [None]\n",
    "\n",
    "    print(name)\n",
    "    print(versions[0])\n",
    "    print(reference_genomes[0] + \"\\n\")\n",
    "\n",
    "    path = [dataset[\"url\"][\"gcp\"][\"us\"]\n",
    "            for dataset in datasets[name][\"versions\"]\n",
    "            if all([dataset[\"version\"] == versions[0],\n",
    "                    dataset[\"reference_genome\"] == reference_genomes[0]])]\n",
    "    assert len(path) == 1\n",
    "    path = path[0]\n",
    "    if path.endswith(\".ht\"):\n",
    "        table = hl.methods.read_table(path)\n",
    "        table_class = \"hail.Table\"\n",
    "    else:\n",
    "        table = hl.methods.read_matrix_table(path)\n",
    "        table_class = \"hail.MatrixTable\"\n",
    "\n",
    "    description = table.describe(handler=lambda x: str(x)).split(\"\\n\")\n",
    "    description = \"\\n\".join([line.rstrip() for line in description])\n",
    "\n",
    "    template = \"\"\".. _{dataset}:\n",
    "\n",
    "{dataset}\n",
    "{underline1}\n",
    "\n",
    "*  **Versions:** {versions}\n",
    "*  **Reference genome builds:** {ref_genomes}\n",
    "*  **Type:** :class:`{class}`\n",
    "\n",
    "Schema ({version0}, {ref_genome0})\n",
    "{underline2}\n",
    "\n",
    ".. code-block:: text\n",
    "\n",
    "{schema}\n",
    "\n",
    "\"\"\"\n",
    "    context = {\n",
    "        \"dataset\": name,\n",
    "        \"underline1\": len(name) * \"=\",\n",
    "        \"version0\": versions[0],\n",
    "        \"ref_genome0\": reference_genomes[0],\n",
    "        \"versions\": \", \".join([str(version) for version in versions]),\n",
    "        \"ref_genomes\": \", \".join([str(reference_genome) for reference_genome in reference_genomes]),\n",
    "        \"underline2\": len(\"\".join([\"Schema (\", str(versions[0]), \", \", str(reference_genomes[0]), \")\"])) * \"~\",\n",
    "        \"schema\": textwrap.indent(description, \"    \"),\n",
    "        \"class\": table_class\n",
    "    }\n",
    "    with open(output_dir + f\"/{name}.rst\", \"w\") as f:\n",
    "        f.write(template.format(**context).strip())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.8"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
